{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用phiGnet预测蛋白功能实操演示\n",
    "\n",
    "## 相关链接\n",
    "\n",
    "- [论文](https://www.nature.com/articles/s41467-024-50955-0)\n",
    "- [源代码:模型+参数](https://doi.org/10.5281/zenodo.12496869),寄存在zenodo网站上,如果连接失败需要修改host文件的ip\n",
    "- [amoai:预测需要使用的进化数据服务器](https://kornmann.bioch.ox.ac.uk/jang/services/amoai/submission.html)\n",
    "- [Grad-CAM keras示例](https://keras.io/examples/vision/grad_cam/)\n",
    "- [神经网络的可解释性（可视化篇）：CAM/Grad-CAM/Grad-CAM++及相关代码（TensorFlow和Pytorch）](https://zhuanlan.zhihu.com/p/479485138)\n",
    "\n",
    "## 源码协议\n",
    "\n",
    "open cc.\n",
    "\n",
    "## 源码依赖\n",
    "\n",
    "源码示例使用的环境:\n",
    "\n",
    "- python==3.6.9\n",
    "- tensorflow-gpu==2.3.1(用pip安装)\n",
    "- numpy==1.18.5\n",
    "\n",
    "要自行生成输入文件,还需要安装pytorch,esm并下载esm1b模型(约15min,7.3GB,自动下载到~/.cache/torch/hub/checkpoints):\n",
    "\n",
    "- conda install pytorch-gpu,pandas\n",
    "- pip install fair-esm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import esm\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本笔记以金葡菌的alpha-溶血素(PDBid:7AHL,UniProtKB:P09616)/肠毒素A(PDBid:1ESF,UniProtKB:P0A0L2)序列为例生成esm1b嵌入:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ahl_seq = 'MKTRIVSSVTTTLLLGSILMNPVAGAADSDINIKTGTTDIGSNTTVKTGDLVTYDKENGMHKKVFYSFIDDKNHNKKLLVIRTKGTIAGQYRVYSEEGANKSGLAWPSAFKVQLQLPDNEVAQISDYYPRNSIDTKEYMSTLTYGFNGNVTGDDTGKIGGLIGANVSIGHTLKYVQPDFKTILESPTDKKVGWKVIFNNMVNQNWGPYDRDSWNPVYGNQLFMKTRNGSMKAADNFLDPNKASSLLSSGFSPDFATVITMDRKASKQQTNIDVIYERVRDDYQLHWTSTNWKGTNTKDKWTDRSSERYKIDWEKEEMTN'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sea_seq = 'MKKTAFTLLLFIALTLTTSPLVNGSEKSEEINEKDLRKKSELQGTALGNLKQIYYYNEKAKTENKESHDQFLQHTILFKGFFTDHSWYNDLLVDFDSKDIVDKYKGKKVDLYGAYYGYQCAGGTPNKTACMYGGVTLHDNNRLTEEKKVPINLWLDGKQNTVPLETVKTNKKNVTVQELDLQARRYLQEKYNLYNSDVFDGKVQRGLIVFHTSTEPSVNYDLFGAQGQYSNTLLRIYRDNKTINSENMHIDIYLYTS'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def esm1b_embed_for_seq(seq:str) -> np.ndarray:\n",
    "    '''生成指定蛋白序列(单字母缩写)的esm1b嵌入'''\n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "    model.eval()  #设置esm_1b模型为预测模式而非训练模式\n",
    "    data = [('_', seq)]\n",
    "    batch_labels, batch_strs, batch_tokens = batch_converter(data)\n",
    "    with torch.no_grad():\n",
    "        results = model(batch_tokens, repr_layers=[33], return_contacts=True)\n",
    "        esm1b_embedding = results['representations'][33]  # 包含esm首尾2个token,长度比蛋白序列长2\n",
    "    return np.squeeze(np.array(esm1b_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ahl_embedding=esm1b_embed_for_seq(ahl_seq)\n",
    "sea_embedding=esm1b_embed_for_seq(sea_seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 源码文件结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "phiGnet_source_dirt='/home/regen/projects/phignet' #源代码解压位置,自行修改"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[01;34mdata\u001b[0m/      \u001b[01;34mmodels\u001b[0m/  predict.py    README.md        \u001b[01;34msrc\u001b[0m/\n",
      "\u001b[01;34mexamples\u001b[0m/  \u001b[01;34moutput\u001b[0m/  \u001b[01;34m__pycache__\u001b[0m/  requirement.txt\n"
     ]
    }
   ],
   "source": [
    "%ls $phiGnet_source_dirt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- data/:存放训练用数据,`collect_data.sh`从指定数据库下载训练需要的GO/EC等数据\n",
    "- examples/:2个示例蛋白序列的输入(npz格式)\n",
    "- models/:模型(hdf5格式)和模型参数(json格式)\n",
    "- output/:预测结果的输出文件夹,初始为空\n",
    "- src/:模型的python代码\n",
    "- predict.py:预测脚本,源码禁用了GPU设备,如使用GPU自行修改\n",
    "\n",
    "## 蛋白功能预测"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 1:下载训练数据(可选)\n",
    "\n",
    "如果想对论文所有结果进行重现,需下载相关数据库文件:\n",
    "\n",
    "```bash\n",
    "bash ./data/collect_data.sh\n",
    "```\n",
    "\n",
    "> 2:生成进化数据\n",
    "\n",
    "给定蛋白序列的进化耦合EVCs和残基群落RCs数据可以在[amoai在线服务](https://kornmann.bioch.ox.ac.uk/jang/services/amoai/index.html)生成.\n",
    "\n",
    "在序列输入框中将蛋白序列或多序列比对进行输入,注意蛋白(或MSA第一个蛋白)长度不要超过500个氨基酸.\n",
    "\n",
    "![amoai服务器提交示例](amoai_example.png)\n",
    "\n",
    "1. 选择`Evolutionary residue communities`\n",
    "2. 输入工作名\n",
    "3. 输入你的名字\n",
    "4. 输入你的邮箱,将发送到该邮箱\n",
    "5. 选择输入类型为单序列还是多序列比对序列\n",
    "6. 输入序列(或上传fasta文件)\n",
    "7. 提交请求,等待约1min,跳出计算进程页面即可关闭(或在该页面等待或者添加书签),共计约20min\n",
    "8. 计算完成后在邮箱页面生成的网址中下载(Figure 4. Inferred evolutionary couplings...+Figure 5. Residue communities旁的链接,csv格式文件,本笔记仓库中重命名为ahl_(evc|rc).csv)\n",
    "9. 打包所有输入为npz文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def phiGnet_input_merge(seq:str,evc_csv:str,rc_csv:str,esm1b_ndarray,jobname:str) -> None:\n",
    "    '''输入序列,evc文件,rc文件,esm1b numpy数组和输出的文件名前缀,在当前目录下生成phiGnet的npz输入文件'''\n",
    "    seq_nd=np.array(seq,np.str_)\n",
    "    evc_df =pd.read_csv(evc_csv)\n",
    "    evc_nd=evc_df.iloc[:, 1:].to_numpy(dtype=np.float32)  #第1列为残基编号,不需要\n",
    "    rc_dummy=evc_df.copy() # rc表格只保存有关联的残基对,需要复制evc表格的尺寸和标签\n",
    "    rc_dummy.iloc[:,1:]=0.0 # 除标签外,全部清空复制表格\n",
    "    rc_df = pd.read_csv(rc_csv)\n",
    "    rc_merge = pd.merge(rc_dummy, rc_df, how='outer', sort=False).fillna(0.0) #合并,重排为邻接矩阵\n",
    "    rc_merge.drop_duplicates('AA',keep='last',inplace=True)\n",
    "    rc_merge.set_index('AA',inplace=True)\n",
    "    rc_merge.sort_index(\n",
    "        key=np.vectorize(lambda name: rc_merge.columns.get_loc(name)),\n",
    "        inplace=True)\n",
    "    rc_merge = rc_merge.to_numpy(dtype=np.float32)\n",
    "    np.savez(f'{jobname}.npz',coupling=evc_nd,esm_1b=esm1b_ndarray,rc=rc_merge,seq=seq_nd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "phiGnet_input_merge(ahl_seq,'ahl_evc.csv','ahl_rc.csv',ahl_embedding,'ahl')\n",
    "phiGnet_input_merge(sea_seq,'sea_evc.csv','sea_rc.csv',sea_embedding,'sea')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "源码的examples/文件夹中有提前准备好的2个蛋白进化数据.\n",
    "\n",
    "> 3:预测蛋白功能\n",
    "\n",
    "使用predict.py进行预测,如果不在源代码根目录下运行,则需要修改源代码:\n",
    "\n",
    "```py\n",
    "# predict.py 15行\n",
    "dirt_param ='./models/' # 源代码,硬编码为根目录下运行\n",
    "dirt_param = f'{os.path.abspath(__file__).rsplit(\"/\",1)[0]}/models/' # 修改的代码,支持在其他目录运行\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL']='3' #忽略错误以外的TF日志"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: predict.py [-h] [-ont {mf,bp,cc,ec}] [-p PATH_INPUTS] [-j JOB_ID]\n",
      "                  [-ct CUT_THRESH] [-d DIRT]\n",
      "\n",
      "optional arguments:\n",
      "  -h, --help            show this help message and exit\n",
      "  -ont {mf,bp,cc,ec}, --ontology {mf,bp,cc,ec}\n",
      "                        the classes of GO terms (default: mf)\n",
      "  -p PATH_INPUTS, --path_inputs PATH_INPUTS\n",
      "                        the path save the inputs of PhiGnet (default:\n",
      "                        ./examples/)\n",
      "  -j JOB_ID, --job_id JOB_ID\n",
      "                        the name of protein whose function is to be predicted\n",
      "                        (default: )\n",
      "  -ct CUT_THRESH, --cut_thresh CUT_THRESH\n",
      "                        the threshold of filtering EVCs and RCs (default: 0.2)\n",
      "  -d DIRT, --dirt DIRT  the path of output (default: )\n"
     ]
    }
   ],
   "source": [
    "!python $phiGnet_source_dirt/predict.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 用法: `python predict.py [-h] [-ont {mf,bp,cc,ec}] [-p PATH_INPUTS] [-j JOB_ID] [-ct CUT_THRESH] [-d DIRT]`\n",
    "- 参数:\n",
    "  - -h,--help 显示帮助信息\n",
    "  - -ont|--ontology {mf,bp,cc,ec}:计算GO术语,包括分子功能mf,生物过程bp,细胞组分cc,酶学委员会编号ec,默认为mf\n",
    "  - -p|--path_inputs PATH_INPUTS:包含PhiGnet输入文件所在的目录,默认为./exapmles\n",
    "  - -j|--job_id JOB_ID:输入文件的文件名(不带拓展名)\n",
    "  - -ct|--cut_thresh CUT_THRESH:过滤EVCs和RCs的阈值,默认0.2\n",
    "  - -d|--dirt DIRT:输出文件的目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont mf -p ./ -j ahl -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "ahl,MF,GO:0005488,0.76840,binding\n"
     ]
    }
   ],
   "source": [
    "%cat ahl_MF_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont bp -p ./ -j ahl -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "ahl,BP,GO:0001906,0.92989,cell killing\n",
      "ahl,BP,GO:0031640,0.79657,killing of cells of other organism\n",
      "ahl,BP,GO:0044419,0.64617,biological process involved in interspecies interaction between organisms\n",
      "ahl,BP,GO:0009987,0.49572,cellular process\n"
     ]
    }
   ],
   "source": [
    "%cat ahl_BP_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont cc -p ./ -j ahl -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "ahl,CC,GO:0005576,1.00000,extracellular region\n",
      "ahl,CC,GO:0110165,1.00000,cellular anatomical entity\n"
     ]
    }
   ],
   "source": [
    "%cat ahl_CC_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont ec -p ./ -j ahl -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,EC_number,Score,Function\n"
     ]
    }
   ],
   "source": [
    "%cat ahl_EC_predicted.csv # 溶血素未能预测EC,可能性过低"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "alpha-溶血素是金葡菌分泌的外毒素,由7条相同的肽链组成,在动物的红细胞膜上组装为孔道蛋白,导致血细胞破裂.以上预测与UniProt上的GO注释基本符合."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont mf -p ./ -j sea -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "sea,MF,GO:0003824,0.94129,catalytic activity\n",
      "sea,MF,GO:0004553,0.68056,\"hydrolase activity, hydrolyzing O-glycosyl compounds\"\n"
     ]
    }
   ],
   "source": [
    "%cat sea_MF_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont bp -p ./ -j sea -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "sea,BP,GO:0044419,0.46557,biological process involved in interspecies interaction between organisms\n"
     ]
    }
   ],
   "source": [
    "%cat sea_BP_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont cc -p ./ -j sea -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,GO_term,Score,Function\n",
      "sea,CC,GO:0005576,1.00000,extracellular region\n",
      "sea,CC,GO:0110165,1.00000,cellular anatomical entity\n"
     ]
    }
   ],
   "source": [
    "%cat sea_CC_predicted.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python $phiGnet_source_dirt/predict.py -ont ec -p ./ -j sea -d ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Predicted by PhiGnet Version 1.0.1\n",
      "ID,Branch,EC_number,Score,Function\n",
      "sea,EC,3.4.21.-,0.79852,3.4.21.-\n"
     ]
    }
   ],
   "source": [
    "%cat sea_EC_predicted.csv # 肠毒素可以预测EC,预测为肽段水解酶,错误"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "肠毒素A也是金葡菌分泌的外毒素,与T细胞受体和2个MHC II直接结合,导致严重的免疫反映.phiGnet预测的BP和CC基本符合,EC和MF预测未见到相关论文."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 预测残基贡献\n",
    "\n",
    "源代码未提供Grad-CAM打分实现,本笔记仓库中的`predict.py`修改为个人猜测的实现,不保证正确性.\n",
    "\n",
    "想要尝试Grad-CAM,可以将本仓库的`predict.py`替换源代码中的脚本,注意备份,然后重新运行上述预测命令,生成额外的残基贡献文件."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'氨基酸的单/三字母缩写映射,未定义的键对应三字母缩写为-'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import defaultdict\n",
    "aa_mono='A F C U D N E Q G H L I K O M P R S T V W Y'\n",
    "aa_tri='Ala Phe Cys Sec Asp Asn Glu Gln Gly His Leu Ile Lys Pyl Met Pro Arg Ser Thr Val Trp Tyr'\n",
    "aa=defaultdict(lambda:'-')\n",
    "aa.update({mono:tri for mono,tri in zip(aa_mono.split(),aa_tri.split())})\n",
    "'''氨基酸的单/三字母缩写映射,未定义的键对应三字母缩写为-'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List,Tuple\n",
    "def best_n_aa(seq:str,contribution_file:str,best_n:int=-1) -> List[Tuple[int,str,float]]:\n",
    "    '''对给定残基序列和贡献文件,输出贡献分数最高的前n个残基信息(编号,氨基酸缩写,贡献分数).\n",
    "\n",
    "    不指定best_n,或者best_n超过序列长度,将返回全体残基的信息.\n",
    "\n",
    "    编号是序列中的自然序号,从1开始,不一定与残基的实验序号一致.\n",
    "\n",
    "    贡献最高的残基分数为100,其余依次下降.\n",
    "    '''\n",
    "    output=[]\n",
    "    if best_n<0:\n",
    "        best_n=len(seq)\n",
    "    cont_vec=np.loadtxt(contribution_file)\n",
    "    seq_tri=[aa[mono] for mono in seq]\n",
    "    for idx,(tri,score) in enumerate(zip(seq_tri,cont_vec),1):\n",
    "        output.append((idx,tri,score*100))\n",
    "    output.sort(key=lambda x:x[2],reverse=True)\n",
    "    output=output[:best_n]\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(163, 'Gly', 100.0),\n",
       " (171, 'Thr', 99.25720691680908),\n",
       " (165, 'Asn', 95.00413537025452),\n",
       " (168, 'Ile', 93.82355809211731),\n",
       " (7, 'Ser', 89.58629369735718),\n",
       " (161, 'Leu', 89.40049409866333),\n",
       " (11, 'Thr', 88.3443832397461),\n",
       " (166, 'Val', 86.71526908874512),\n",
       " (21, 'Asn', 86.70112490653992),\n",
       " (160, 'Gly', 83.29342007637024),\n",
       " (13, 'Leu', 82.7286422252655),\n",
       " (152, 'Gly', 82.58850574493408),\n",
       " (5, 'Ile', 82.54868388175964),\n",
       " (10, 'Thr', 81.9492518901825),\n",
       " (170, 'His', 78.79922389984131)]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_n_aa(ahl_seq,'ahl_gradcam_bp_GO:0001906_cell_killing_predicted.txt',15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(17, 'Ser', 100.0),\n",
       " (18, 'Ile', 88.89694213867188),\n",
       " (13, 'Leu', 78.43815088272095),\n",
       " (11, 'Thr', 64.62135314941406),\n",
       " (21, 'Asn', 39.60690498352051),\n",
       " (12, 'Thr', 38.700222969055176),\n",
       " (19, 'Leu', 38.613638281822205),\n",
       " (6, 'Val', 32.45473504066467),\n",
       " (20, 'Met', 31.147804856300354),\n",
       " (8, 'Ser', 27.847659587860107),\n",
       " (16, 'Gly', 26.292625069618225),\n",
       " (186, 'Pro', 25.346654653549194),\n",
       " (7, 'Ser', 24.30262118577957),\n",
       " (313, 'Glu', 24.2129385471344),\n",
       " (209, 'Asp', 23.885059356689453)]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_n_aa(ahl_seq,'ahl_gradcam_mf_GO:0005488_binding_predicted.txt',15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3D可视化\n",
    "\n",
    "使用pymol进行,由于conda编译的pymol不支持python3.6,本笔记仅生成pymol脚本,可视化在其他环境进行.\n",
    "\n",
    "根据每个残基的贡献分数对蛋白表面进行着色.\n",
    "\n",
    "可选方案:使用同源建模或生成方法(alphafold),以pdb3D结构为模板,预测完整序列的3D结构,在该结构上着色,以下脚本可简化,仅适合孤立的蛋白.涉及的工具太多,本笔记不另行实现."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def py_for_PhiGnet_aa(contribute_list: List[Tuple[int,str,float]], pdbid: str = '',\n",
    "                      local_file: str = '', protein_chain: str = 'A', aa_offset: int = 0) -> None:\n",
    "    '''根据PhiGnet预测的残基贡献生成pymol脚本文件,对蛋白表面进行着色.在当前目录保存到`(pdbid|local_file).py`文件.\n",
    "\n",
    "    pdbid:4位id如7ahl(大小写不敏感).\n",
    "\n",
    "    指定local_file时,将忽略pdbid,从本地加载结构文件.\n",
    "\n",
    "    protein_chain指定要强调的链号,以白色不透明表面表示.\n",
    "\n",
    "    aa_offset是PhiGnet输入序列残基编号-3D结构残基编号的偏移量,避免实验结构残基编号与序列编号不一致.函数假定实验结构残基序列是连续的,可能有些测定不良的结构会出问题.\n",
    "    '''\n",
    "\n",
    "    contribute_list = contribute_list.copy()\n",
    "    contribute_list.sort()\n",
    "\n",
    "    if aa_offset:\n",
    "        contribute_list = [(idx - aa_offset, resname, score)\n",
    "                           for idx, resname, score in contribute_list if idx- aa_offset>0]#掐头,如果序列尾部比实验结构长,pymol选择语法会自动忽略\n",
    "\n",
    "    if not pdbid and not local_file:\n",
    "        raise ValueError('至少需要提供pdbid或local_file的参数')\n",
    "    if local_file:\n",
    "        output_name=local_file.rsplit(('/','\\\\'),-1)+'.predicted.py'\n",
    "    else:\n",
    "        output_name=pdbid+'.predicted.py'\n",
    "    p_chain_selection = f'chain {protein_chain}'\n",
    "\n",
    "    color_template=[]\n",
    "    for idx, resname, score in contribute_list:\n",
    "        if score>0:\n",
    "            color_template.append(f'cmd.set_color(\"phiGnet{idx}\", {[255, 255-int(255 * score / 100), 255-int(255 * score / 100)]})')\n",
    "            color_template.append(f'''cmd.color(\"phiGnet{idx}\", \"{p_chain_selection}{'' if resname=='-' else ' and resn '+resname} and resi {idx}\")''')\n",
    "    color_template='\\n'.join(color_template)\n",
    "\n",
    "    template = f'''import pymol.cmd as cmd\n",
    "cmd.{'load' if local_file else 'fetch'}('{local_file if local_file else pdbid}', 'PhiGnet_focus')\n",
    "cmd.remove('solvent')\n",
    "\n",
    "p_chain_selection = 'chain {protein_chain}'\n",
    "o_chain_selection = 'not (chain {protein_chain})'\n",
    "\n",
    "cmd.show_as('surface', o_chain_selection)\n",
    "cmd.color('blue', o_chain_selection)\n",
    "cmd.set('transparency', 0.7, o_chain_selection)\n",
    "\n",
    "cmd.color('white', p_chain_selection)\n",
    "cmd.show('surface', p_chain_selection)\n",
    "cmd.set('transparency', 0.2, p_chain_selection)\n",
    "\n",
    "{color_template}\n",
    "'''\n",
    "\n",
    "    with open(f'{output_name}.py','w+',encoding='utf-8') as outputfile:\n",
    "        outputfile.write(template)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在当前目录下,用可运行pymol的环境进行绘图.命令行`pymol xxxxx.py`或者在图形界面>File>Open打开."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "py_for_PhiGnet_aa(best_n_aa(ahl_seq,'ahl_gradcam_bp_GO:0001906_cell_killing_predicted.txt'),'7ahl',aa_offset=26) \n",
    "#以alpha-溶血素bp中得分最高的细胞损伤能力为例.实验结构第1个序列为完整序列的第27个,偏移量为27-1=26"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alpha-溶血素参与cell killing过程的重要残基](ahl_bp_cell_killing.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "py_for_PhiGnet_aa(best_n_aa(ahl_seq,'ahl_gradcam_mf_GO:0005488_binding_predicted.txt'),'7ahl',aa_offset=26) #预测的结果分数最高的集中在N端,在3D图像中未着色"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alpha-溶血素自组装的关键位点](ahl_mf_binding.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 个人疑问\n",
    "\n",
    "esm-1b嵌入的向量长度=序列长度+2,esm-1b/evcs/rcs/序列独热编码长度均补全到1024,那么esm-1b和序列不就差一位了吗?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 清理生成的文件\n",
    "\n",
    "%rm *predicted*\n",
    "%rm *.cif\n",
    "%rm *.npz"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cctbx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
